#ifndef LQR_H_
#define LQR_H_

#include <armadillo>
#include "matfunction.h"

#define TIME_INVARIANT_FINITE 			0
#define TIME_INVARIANT_INFINITE 		1
#define TIME_VARIANT_FINITE  			2


class lqrSolver
{
private:
	/* parameters of the system */

	// dimension of the system;
	int dim; 

	int type;

	// time-variant case
	MatFunction A_t;
	MatFunction B_t;
	MatFunction S_t;
	MatFunction Q_t;
	MatFunction R_t;

	// discrete expression
	std::vector<arma::mat> An;
	std::vector<arma::mat> Bn;
	std::vector<arma::mat> Qn;
	std::vector<arma::mat> Rn;
	std::vector<arma::mat> Hn;

	// time-invariant case
	arma::mat A;
	arma::mat B;
	arma::mat S;
	arma::mat Q;
	arma::mat R;

	// time interval
	double t0;
	double t1;

	int length;

	// reshaped RDE
	// \dot{\mat{X \\ Y}} = \mat{H}\mat{X \\ Y}
	arma::mat H; // initialized automatically
	std::vector<arma::mat> XYn;
	std::vector<arma::mat> Pn;

	// boundary conditions
	arma::mat In;
	arma::mat x0;

	double step;

	// data
	std::vector<arma::mat> u;
	std::vector<arma::mat> x;
	
	// auto-initialization
	void init();

	// update An, Bn, Qn, Rn and Hn
	void update(int n);

	// default line scheme
	const char* lineScheme[10] = {"p-1", "r|1", "g;1", "b:1", "cj1", "mi1", "y=1", "l1", "e1", "k1"};
public:
	// set the parameters

	void setType(int _type)
	{
		type = _type;
	};

	void setA_t(MatFunction _A_t);
	void setB_t(MatFunction _B_t);
	void setS_t(MatFunction _S_t);
	void setQ_t(MatFunction _Q_t);
	void setR_t(MatFunction _R_t);

	void setA(arma::mat _A){ A = _A; };
	void setB(arma::mat _B){ B = _B; };
	void setS(arma::mat _S){ S = _S; };
	void setQ(arma::mat _Q){ Q = _Q; };
	void setR(arma::mat _R){ R = _R; };

	void setTime(double _t0, double _t1){t0 = _t0; t1 = _t1;};
	void setX0(arma::mat _x0){ x0 = _x0;};
	void setStep(double m_step){ step = m_step;};
	
	lqrSolver(	arma::mat A, 
			arma::mat B, 
			arma::mat S, 
			arma::mat Q, 
			arma::mat R, 
			double t0, 
			double t1, 
			arma::mat x0, 
			double m_step)
	{
		setA(A);
		setB(B);		
		setS(S);
		setQ(Q);
		setR(R);
		setTime(t0, t1);
		setX0(x0);
		setStep(m_step);
		type = TIME_INVARIANT_FINITE;
	};
	
	lqrSolver(	arma::mat A, 
			arma::mat B, 
			arma::mat S, 
			arma::mat Q, 
			arma::mat R, 
			double t0, 
			double t1,
			arma::mat x0)
	{
		setA(A);
		setB(B);		
		setS(S);
		setQ(Q);
		setR(R);
		setTime(t0, t1);
		setX0(x0);
		step = 0.01;
		type = TIME_INVARIANT_FINITE;
		
	};

	lqrSolver(	arma::mat A, 
			arma::mat B,  
			arma::mat Q, 
			arma::mat R, 
			double t0, 
			double t1, 
			arma::mat x0, 
			double m_step)
	{
		setA(A);
		setB(B);		
		setQ(Q);
		setR(R);
		setTime(t0, t1);
		setX0(x0);
		setStep(m_step);
		type = TIME_INVARIANT_INFINITE;
	};

	lqrSolver(	arma::mat A, 
			arma::mat B,  
			arma::mat Q, 
			arma::mat R, 
			double t0, 
			double t1,
			arma::mat x0)
	{
		setA(A);
		setB(B);		
		setQ(Q);
		setR(R);
		setTime(t0, t1);
		setX0(x0);
		step = 0.01;
		type = TIME_INVARIANT_INFINITE;
		
	};

	lqrSolver(	MatFunction _A_t, 
			MatFunction _B_t, 
			MatFunction _S_t, 
			MatFunction _Q_t, 
			MatFunction _R_t, 
			double t0, 
			double t1, 
			arma::mat x0, 
			double m_step)
	{
		setA_t(_A_t);
		setB_t(_B_t);		
		setS_t(_S_t);
		setQ_t(_Q_t);
		setR_t(_R_t);
		setTime(t0, t1);
		setX0(x0);
		setStep(m_step);
		type = TIME_VARIANT_FINITE;
	};

	// solve the RDE by Sympletic method
	void solve();

	// save the result to file
	void saveData(const char* filename ) const;
	void saveData() const;

	// draw the result with MathGL
	void draw() const;
	void draw(const char* filename_prefix) const;
	void draw_x(const char* filename) const;
	void draw_x() const;
	void draw_u(const char* filename) const;
	void draw_u() const; 
	void draw_P(const char* filename) const;
	void draw_P() const;	
};


#endif
